from turtle import forward
from torch import relu
import torch
import torch.nn as nn
import torch.nn.functional as F

class LeNet(nn.Module):
    # 参数化激活函数后，可以切换不同的激活函数如Sigmoid,GELU,Tannh,ELU
    def __init__(self,activation_fn):
        super().__init__()
        
        """
        分别对应的含义: (in_channels:输入通道数, out_channels:输出通道数, 
        kernel_size:卷积核的大小并不是数量,stride =1: 步长默认为1可以省略,padding = 0: 默认为0)
        计算公式: 经过卷积后的矩阵尺寸大小计算公式为 N = (M - F + 2P)/S +1
        其中 输入图片大小为MxM,Filter大小为FxF,步长s,padding 为p
        """
        self.net = nn.Sequential(
            nn.Conv2d(3, 6, 5, 1, 0),activation_fn,    # input(3,32,32)   output(6,28,28)
            nn.MaxPool2d(2, 2, 0),                 # output(6,14,14)
            nn.Conv2d(6, 16, 5, 1, 0),activation_fn,   # output(16,10,10)
            nn.MaxPool2d(2, 2, 0),                 # output(16,5,5)            
            nn.Flatten(),
            nn.Linear(16*5*5, 120),activation_fn,
            nn.Linear(120, 84),activation_fn, 
            nn.Linear(84, 10)
        )
        
    def forward(self,x):
        return self.net(x)
